import { MongoClient } from "mongodb";
import { GovDocTextSplitter } from "./docParseHelper.js";
import {
  MONGO_URI,
  MONGO_DB,
  SOURCE_COLLECTION,
  DESTINATION_COLLECTION,
} from "./config.js";

async function transferData() {
  const client = new MongoClient(MONGO_URI, {
    useNewUrlParser: true,
    useUnifiedTopology: true,
  });

  try {
    // 连接到 MongoDB
    await client.connect();
    const db = client.db(MONGO_DB);
    const sourceCollection = db.collection(SOURCE_COLLECTION);
    const destinationCollection = db.collection(DESTINATION_COLLECTION);
    // 创建一个游标来遍历集合
    // 过滤掉"如来写作" 仅查询以inner开头的
    const cursor = sourceCollection.find({
      // platform_id: { $ne: "rulaixiezuo" },
      platform_id: "inner_qstheory_hongqi"
    });
    // 批量处理数据
    const batchSize = 1000; // 每次处理的批量大小
    let batch = [];

    while (await cursor.hasNext()) {
      const doc = await cursor.next();
      const {
        _id,
        id,
        platform_id,
        category_type,
        origin_category,
        title,
        url,
        content,
      } = doc;
      const article_id = id;
      const document_id = _id.toString();
      //   如果小于 50 文章过短  忽略
      if (content.length < 100) {
        continue;
      }
      // 去除标题
      const cleanContent = content.replace(title, "").trim();
      const ariticleChunks = GovDocTextSplitter.splitText(cleanContent);
      let len = ariticleChunks.length;
      // 如果为空忽略
      if (len === 0) {
        continue;
      }
      // 检查最后一个片段 如果太短 则与前一个合并 （合并前数据库：37262）
      if (len > 1 && ariticleChunks[len - 1].trim().length < 100) {
        ariticleChunks[len - 2] += ariticleChunks[len - 1];
        ariticleChunks.pop();
        len--;
      }
      console.log(`Processing article: ${url} -- ${title}`);
      for (let i = 0; i < len; i++) {
        // console.log(`第${i}片长度${cleanedChunk.length}==》${cleanedChunk}`);
        const cleanedChunk = ariticleChunks[i].trim();
        const fragmentId = article_id + `_${i}`;
        // 检查是否已存在
        const exists = await destinationCollection.findOne({ id: fragmentId });
        if (exists) {
          console.log(`fragment with id=${fragmentId} already exists`);
          continue;
        }
        const fragment = {
          id: fragmentId,
          document_id: document_id, // 文章ID
          fragment_index: i,
          fragment_content: cleanedChunk,
          platform_id,
          category_type,
          origin_category,
          title,
          url,
        };
        batch.push(fragment);
      }
      // 批量插入 每次插入
      if (batch.length >= batchSize) {
        try {
          await destinationCollection.insertMany(batch, { ordered: false });
          console.log(`Inserted batch of ${batch.length} documents`);
        } catch (err) {
          console.error(`Error inserting batch: ${err}`);
        }
        batch = [];
      }
    }
    // 插入剩余的文档
    if (batch.length > 0) {
      try {
        await destinationCollection.insertMany(batch, { ordered: false });
        console.log(`Inserted final batch of ${batch.length} documents`);
      } catch (err) {
        console.error(`Error inserting final batch: ${err}`);
      }
    }
  } catch (err) {
    console.error(err);
  } finally {
    // 关闭 MongoDB 连接
    await client.close();
  }
}

// 执行数据转移
transferData().catch(console.error);
